import React from 'react';
import { Table, TableProps } from 'antd';
import { DndContext, DragEndEvent } from '@dnd-kit/core';
import { restrictToVerticalAxis } from '@dnd-kit/modifiers';
import { SortableContext, verticalListSortingStrategy } from '@dnd-kit/sortable';
import DragTableRow from '../DragTableRow';

interface DragTableProps<RecordType, KeyType extends React.Key = React.Key> 
  extends Omit<TableProps<RecordType>, 'rowKey'> {
  /**
   * 拖拽结束的回调函数
   */
  handleDragEnd?: (event: DragEndEvent) => void;
  
  /**
   * 获取记录唯一键的函数
   * @default (record, index) => record.id ?? index
   */
  getRowKey?: (record: RecordType, index: number) => KeyType;
  
  /**
   * 是否启用拖拽
   * @default true
   */
  draggable?: boolean;
}

function DragTable<RecordType>(props: DragTableProps<RecordType>) {
  const { 
    handleDragEnd, 
    getRowKey = (record, index) => (record as any).id ?? index,
    draggable = true,
    ...tableProps 
  } = props;

  const rowKey = (record: RecordType, index?: number) => 
    getRowKey(record, index!);

  return draggable ? (
    <DndContext modifiers={[restrictToVerticalAxis]} onDragEnd={handleDragEnd}>
      <SortableContext
        items={tableProps.dataSource?.map((record, index) => 
          getRowKey(record, index)) || []}
        strategy={verticalListSortingStrategy}
      >
        <Table<RecordType>
          {...tableProps}
          rowKey={rowKey}
          components={{
            ...tableProps.components,
            body: {
              ...tableProps.components?.body,
              row:DragTableRow
              // 这里可以添加你的自定义行组件
            }
          }}

        />
      </SortableContext>
    </DndContext>
  ) : (
    <Table<RecordType>
      {...tableProps}
      rowKey={rowKey}
    />
  );
}

export default DragTable;